Refactor BaseTableDataManager by krishan1390 · Pull Request #18381 · apache/pinot

krishan1390 · 2026-04-30T06:54:54Z

Summary

Small refactors to BaseTableDataManager (and its tests) to enable cleaner extensions, abstractions, and ownership for future work.

1. Push `_segmentReloadSemaphore` acquire/release into `reloadSegment(SegmentDataManager, ILC, boolean)`

Removes the duplicated outer acquire/release in the reloadSegment(String) and parallel reloadSegments(List<SDM>) paths — there is now a single acquisition site.

2. Centralize segment delete in `BaseTableDataManager`

Segment deletion belongs on the TDM, not on HelixInstanceDataManager. This change moves it there:

New SPI method: TableDataManager#deleteSegment(String).

3. Drop `destroy()`-without-`offload()` assertion in `SegmentDataManager`

destroy() already calls offload() first, so callers can simply do:

if (segmentDataManager.decreaseReferenceCount()) {
    segmentDataManager.destroy();
}

instead of the longer pattern that requires an explicit prior offload(). The previous assertion forced the longer pattern with no real benefit.

4. Split `tryLoadExistingSegment` so subclasses can load without registering

Introduces tryLoadExistingSegmentInternal which returns the ImmutableSegment (or null on miss/stale CRC/load failure) without invoking addSegment or other registration hooks. The public tryLoadExistingSegment is now a thin wrapper that calls the internal method and registers the result.

5. Test refactor for extensibility

BaseTableDataManagerAcquireSegmentTest is reworked into a reusable base so subclasses can run all existing test bodies against a different TableDataManager implementation. Exposed hooks:

newTableDataManager() — choose the concrete TDM under test.
tdmKey(segmentName) — map a logical segment name to the TDM-map key.
addSegment(tdm, seg), getInnerSegment(sdm) — small indirections so subclasses can vary segment registration / wrapping.

Adjacent visibility tweaks (isSegmentStale → protected, hasSameCRC → public) support the same goal.

SPI changes

New: TableDataManager#deleteSegment(String).

Test plan

Existing BaseTableDataManager and subclass tests pass.
Reload / segment-load / delete paths exercised in integration tests.

…ts building blocks Three small refactors that keep single-segment behavior identical and let a future multi-segment SegmentDataManager (e.g. a wrapper around N constituent segments) reuse the same load and reload primitives without forking them: 1. Extract `protected ImmutableSegment loadSegment(zkMetadata, ilc)` from `downloadAndLoadSegment`. The new helper performs only the download + `ImmutableSegmentLoader.load`, returning the segment without registering it in `_segmentDataManagerMap` or invoking upsert hooks. Single-segment callers continue to use `downloadAndLoadSegment`, which now composes the helper + `addSegment(...)`. This lets a multi-segment manager load all of its members first and register a single wrapper entry under one name. 2. Push `_segmentReloadSemaphore` acquire/release down into `reloadSegment(SegmentDataManager, IndexLoadingConfig, boolean)`. The public `reloadSegment(String)` and the private parallel `reloadSegments(List<SDM>)` both used to wrap the inner call with the semaphore; that acquire is now inside the per-physical-segment body and the outer wrappers are removed (which would otherwise double-acquire on a non-reentrant semaphore). For non-group tables this is structurally identical (one segment -> one acquire -> one release; same concurrency bound). For multi-segment managers that fan out N reloads, each member contends for a slot independently. 3. Drop `@VisibleForTesting` on `isSegmentStale(IndexLoadingConfig, SegmentDataManager)` and widen to plain `protected` so subclasses can call it from group-aware overrides of `getStaleSegments` / `needReloadSegments`. The semaphore stays at the orchestration boundary in `doReplaceSegment` (not relocated into `replaceSegmentIfCrcMismatch`), because subclasses commonly override `replaceSegmentIfCrcMismatch` and a relocation there would leak the acquire across paths that bypass the override; subclasses needing per-member acquire on a multi-segment replace can wrap the call themselves. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

codecov-commenter · 2026-04-30T07:44:36Z

Codecov Report

❌ Patch coverage is 57.89474% with 24 lines in your changes missing coverage. Please review.
✅ Project coverage is 63.67%. Comparing base (bd76ff0) to head (25af3f2).

Files with missing lines	Patch %	Lines
.../pinot/core/data/manager/BaseTableDataManager.java	63.46%	17 Missing and 2 partials ⚠️
...server/starter/helix/HelixInstanceDataManager.java	0.00%	5 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #18381      +/-   ##
============================================
- Coverage     63.68%   63.67%   -0.01%     
  Complexity     1685     1685              
============================================
  Files          3266     3266              
  Lines        199821   199849      +28     
  Branches      31022    31024       +2     
============================================
+ Hits         127257   127258       +1     
- Misses        62425    62450      +25     
- Partials      10139    10141       +2

Flag	Coverage Δ
custom-integration1	`100.00% <ø> (ø)`
integration	`100.00% <ø> (ø)`
integration1	`100.00% <ø> (ø)`
integration2	`0.00% <ø> (ø)`
java-21	`63.67% <57.89%> (-0.01%)`	⬇️
temurin	`63.67% <57.89%> (-0.01%)`	⬇️
unittests	`63.67% <57.89%> (-0.01%)`	⬇️
unittests1	`55.81% <21.15%> (-0.02%)`	⬇️
unittests2	`34.95% <43.85%> (+<0.01%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

…ions and ownership for future work

krishan1390 · 2026-05-12T11:47:51Z

  }

+  @Override
+  public void deleteSegment(String segmentName)


This is moved from HelixInstanceDataManager to TableDataManager as this is the right place

krishan1390 · 2026-05-12T11:54:50Z

-      } finally {
-        _segmentReloadSemaphore.release();
-      }
+      reloadSegment(segmentDataManager, indexLoadingConfig, forceDownload);


moved the sempahore to inside reloadSegment

krishan1390 · 2026-05-12T12:11:29Z

   */
  public void destroy() {
-    // NOTE: We want the test to catch the case when destroy is called without offloading, but not fail the production.
-    assert _offloaded.get() : "Cannot destroy segment data manager without offloading it first";


I think it makes sense for destroy to call offload first if we want the segment to be offloaded only after its reference count becomes 0

Also I don't see a reason to complicate clients to do below loop always

segmentDataManager.offload(); if (segmentDataManager.decreaseReferenceCount()) { segmentDataManager.destroy(); }

Rather below is simpler

if (segmentDataManager.decreaseReferenceCount()) { segmentDataManager.destroy(); }

Thus removing this assertion

krishan1390 · 2026-05-12T12:46:45Z

  private final Set<SegmentDataManager> _accessedSegManagers = ConcurrentHashMap.newKeySet();
  private final Set<SegmentDataManager> _allSegManagers = ConcurrentHashMap.newKeySet();
-  private Map<String, ImmutableSegmentDataManager> _internalSegMap;
+  protected Map<String, SegmentDataManager> _internalSegMap;


These changes are refactoring to enable a more generic and extensible BaseTableDataManagerTest

krishan1390 · 2026-05-13T05:40:49Z

        } else {
          _logger.info("Reloading (force committing) consuming segment: {}", segmentName);
-          ((RealtimeSegmentDataManager) segmentDataManager).forceCommit();
+          _segmentReloadSemaphore.acquire(segmentName, _logger);


Since we were capturing the semaphore earlier too for force commit, we're doing it now too.

krishan1390 and others added 2 commits April 30, 2026 12:23

Cleanup comment

cc2afba

krishan1390 marked this pull request as draft May 4, 2026 12:23

krishan1390 added 2 commits May 6, 2026 13:44

Refactor BaseTDM and its tests to enable cleaner extensions, abstract…

ab53cb5

…ions and ownership for future work

Merge remote-tracking branch 'upstream/master' into basetdm_refactor

1ae36c4

krishan1390 changed the title ~~Refactor BaseTableDataManager so multi-segment managers can compose its building blocks~~ Refactor BaseTableDataManager to enable cleaner extensions, abstractions and ownership for future work May 6, 2026

krishan1390 changed the title ~~Refactor BaseTableDataManager to enable cleaner extensions, abstractions and ownership for future work~~ Refactor BaseTableDataManager May 6, 2026

krishan1390 added 5 commits May 6, 2026 17:31

Cleanup deleteSegment

0eb9597

Move offloadSegment to deleteSegment

953952b

Reload should take lock first and then acquire semaphore

c818a8c

Cleanup BaseTDM

ccdfbb4

Merge remote-tracking branch 'upstream/master' into basetdm_refactor

52fa8e5

krishan1390 commented May 12, 2026

View reviewed changes

krishan1390 added 2 commits May 13, 2026 10:44

Acquire then take segment lock in reload

9a9c567

Acquire reload semaphore for realtime force commit too

3ea1b9c

krishan1390 commented May 13, 2026

View reviewed changes

Make some accessors public in TDM

23afba3

krishan1390 marked this pull request as ready for review May 13, 2026 09:13

krishan1390 added 3 commits May 13, 2026 14:43

Merge remote-tracking branch 'upstream/master' into basetdm_refactor

eb0602b

Make some accessors public in TDM

1f67c75

Merge remote-tracking branch 'upstream/master' into basetdm_refactor

25af3f2

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor BaseTableDataManager #18381

Refactor BaseTableDataManager #18381
krishan1390 wants to merge 15 commits into
apache:masterfrom
krishan1390:basetdm_refactor

krishan1390 commented Apr 30, 2026 •

edited

Loading

Uh oh!

codecov-commenter commented Apr 30, 2026 •

edited

Loading

Uh oh!

krishan1390 May 12, 2026

Uh oh!

krishan1390 May 12, 2026

Uh oh!

krishan1390 May 12, 2026

Uh oh!

Uh oh!

krishan1390 May 12, 2026

Uh oh!

krishan1390 May 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

krishan1390 commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

1. Push _segmentReloadSemaphore acquire/release into reloadSegment(SegmentDataManager, ILC, boolean)

2. Centralize segment delete in BaseTableDataManager

3. Drop destroy()-without-offload() assertion in SegmentDataManager

4. Split tryLoadExistingSegment so subclasses can load without registering

5. Test refactor for extensibility

SPI changes

Test plan

Uh oh!

codecov-commenter commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

krishan1390 May 12, 2026

Choose a reason for hiding this comment

Uh oh!

krishan1390 May 12, 2026

Choose a reason for hiding this comment

Uh oh!

krishan1390 May 12, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

krishan1390 May 12, 2026

Choose a reason for hiding this comment

Uh oh!

krishan1390 May 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

krishan1390 commented Apr 30, 2026 •

edited

Loading

1. Push `_segmentReloadSemaphore` acquire/release into `reloadSegment(SegmentDataManager, ILC, boolean)`

2. Centralize segment delete in `BaseTableDataManager`

3. Drop `destroy()`-without-`offload()` assertion in `SegmentDataManager`

4. Split `tryLoadExistingSegment` so subclasses can load without registering

codecov-commenter commented Apr 30, 2026 •

edited

Loading